{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8aad0f08",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from image_process.EyeDataset import EyeDataset\n",
    "from image_process.noise_manage import events_to_count_images, create_noise_mask, apply_noise_mask\n",
    "from image_process.eyelid_glint import create_eyelid_glint_mask\n",
    "from image_process.eyelash import create_eyelash_mask\n",
    "from image_process.pupil_center import locate_pupil_center, create_pupil_mask\n",
    "from image_process.ellipse_fit import fit_ellipse, ellipse_fit_score\n",
    "\n",
    "data_dir = '.\\\\eye_data'\n",
    "subject = 9\n",
    "eye = 'left'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d0e22643",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Showing the left eye of subject 9\n",
      "\n",
      "Loading Data from .\\eye_data..... \n",
      "\n",
      "Loading Frames....\n",
      "There are 7414 frames \n",
      "\n",
      "Loading Events....\n",
      "There are 123411692 events \n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "生成计数图: 100%|██████████| 15427/15427 [00:16<00:00, 917.27it/s] \n"
     ]
    }
   ],
   "source": [
    "'''把事件转换为极性图片'''\n",
    "\n",
    "eye_dataset = EyeDataset(data_dir, subject)\n",
    "if eye == 'left':\n",
    "    print('Showing the left eye of subject ' + str(subject) + '\\n')\n",
    "    print('Loading Data from ' + data_dir + '..... \\n')\n",
    "    eye_dataset.collect_data(0)\n",
    "else:\n",
    "    print('Showing the right eye of subject ' + str(subject)+ '\\n')\n",
    "    print('Loading Data from ' + data_dir + '..... \\n')\n",
    "    eye_dataset.collect_data(1)\n",
    "    \n",
    "event_list = eye_dataset.event_stack[::-1]\n",
    "\n",
    "images = events_to_count_images(event_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "98807e54",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "保存rgb图像: 100%|██████████| 15427/15427 [02:44<00:00, 93.91it/s] \n"
     ]
    }
   ],
   "source": [
    "'''保存rgb图片'''\n",
    "\n",
    "def pos_neg_to_rgb(pos_img, neg_img):\n",
    "    rgb_img = np.ones(pos_img.shape + (3,), dtype=np.uint8) * 255\n",
    "    rgb_img [(pos_img > 0) & (neg_img == 0)] = [0, 255, 0]  # 绿色\n",
    "    rgb_img [(neg_img > 0) & (pos_img <= 0)] = [255, 0, 0]  # 红色\n",
    "    rgb_img [(pos_img > 0) & (neg_img > 0)] = [255, 255, 0]  # 黄色\n",
    "    return rgb_img\n",
    "\n",
    "rgb_dir = './eye_data/rgb_img'\n",
    "os.makedirs(rgb_dir, exist_ok=True)\n",
    "for idx, (pos_img, neg_img) in enumerate(tqdm(images, desc=\"保存rgb图像\")):\n",
    "    # noise mask\n",
    "    noise_mask = create_noise_mask(pos_img, neg_img, kernel_size=5, thresh_noise=5)\n",
    "    pos_img_denoised, neg_img_denoised = apply_noise_mask(pos_img, neg_img, noise_mask)\n",
    "    # eyelid and glint mask\n",
    "    eyelid_glint_mask = create_eyelid_glint_mask(pos_img_denoised, neg_img_denoised, \n",
    "                         bin_thresh=1, dilate_kernel=7, dilate_shape='ellipse', \n",
    "                         expand_kernel=5, expand_shape='ellipse')\n",
    "    # eyelash mask\n",
    "    eyelash_mask = create_eyelash_mask(pos_img + neg_img, eyelid_glint_mask, \n",
    "                                   blur_ksize=3, blur_thresh=1,\n",
    "                                   morph_ksize1=11, morph_ksize2=5, morph_rect_w=21, morph_rect_h=9)\n",
    "    # iris mask\n",
    "    iris_mask = (noise_mask | eyelid_glint_mask | eyelash_mask)\n",
    "    pos_img_iris, neg_img_iris = apply_noise_mask(pos_img, neg_img, iris_mask)\n",
    "    iris_img = (pos_img_iris + neg_img_iris) > 0\n",
    "\n",
    "    # pupil mask\n",
    "    bandwidth = 32\n",
    "    center_y, center_x = locate_pupil_center(iris_img, bandwidth)\n",
    "    pupil_mask = create_pupil_mask(iris_img, (center_x, center_y), bandwidth)\n",
    "    pos_img_pupil, neg_img_pupil = apply_noise_mask(pos_img_iris, neg_img_iris, ~pupil_mask)\n",
    "\n",
    "    pupil_img = (pos_img_pupil + neg_img_pupil) > 0\n",
    "    points_y, points_x = np.where(pupil_img)\n",
    "    pupil_points = np.column_stack((points_x, points_y))\n",
    "\n",
    "    # ellipse fit and score\n",
    "    ellipse = fit_ellipse(pupil_points)\n",
    "    if ellipse is not None:\n",
    "        fit_score = ellipse_fit_score(pupil_points, ellipse)\n",
    "\n",
    "    # visualize\n",
    "    if ellipse is not None and fit_score > 0.9:\n",
    "        rgb_img = pos_neg_to_rgb(pos_img, neg_img)\n",
    "        rgb_img_pupil = pos_neg_to_rgb(pos_img_pupil, neg_img_pupil)\n",
    "\n",
    "        (x_center, y_center), (width, height), angle = ellipse\n",
    "        center = (int(round(x_center)), int(round(y_center)))\n",
    "        axes = (int(round(width/2)), int(round(height/2)))\n",
    "        cv2.ellipse(rgb_img_pupil, center, axes, angle, 0, 360, (0, 0, 255), 1)\n",
    "        cv2.circle(rgb_img_pupil, center, 1, (0, 0, 255), -1)\n",
    "\n",
    "        combined_img = np.vstack((rgb_img, rgb_img_pupil))\n",
    "        combined_img_bgr = cv2.cvtColor(combined_img, cv2.COLOR_RGB2BGR)\n",
    "        cv2.imwrite(os.path.join(rgb_dir, f'{idx}_{fit_score:.3f}.png'), combined_img_bgr)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "event_camera",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
